# -*- coding: gb2312 -*-

from pymongo import MongoClient
from WindPy import *

import time
import math

class StockGeneralData:
    #include all the stock data information writing and updating, e.g., stock trade status and ipo date
    def __init__(self, baseStartDateStr,  updateDateStr = time.strftime('%Y-%m-%d'), fillType='F', periodStr = 'M', DBIP='localhost', DBPort=27017, DBName = 'stocks', periodSampleType = 'END', stockExtractGrpNum = 10, colIpoDateName = 'ipo_date', colTradeStatusName = 'trade_status', colSysData = 'sys_data'):
        self.baseStartDateStr = baseStartDateStr
        self.updateDateStr = updateDateStr
        self.fillType = fillType
        self.periodStr = periodStr
        self.DBIP = DBIP
        self.DBPort = DBPort
        self.DBName = DBName
        self.periodSampleType = periodSampleType
        self.stockIdList = []
        self.periodTransDateList = []
        self.colIpoDateName = colIpoDateName
        self.colTradeStatusName = colTradeStatusName
        self.colSysData = colSysData
        self.stockExtractGrpNum = stockExtractGrpNum
        w.start()


    def writeBaseIpoDateToMongoDB(self, endDateStr):

        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colName = self.colIpoDateName


        resStockID = w.wset('sectorconstituent', 'date=' + endDateStr + ';sectorid=a001010100000000;field=wind_code')

        windStockIdList = resStockID.Data[0]

        ipoDateDictList = []

        allStockIdStr = ''

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]

        collection_ipo_date = db[colName]

        for i in range(0, len(windStockIdList)):
            allStockIdStr = allStockIdStr + str(windStockIdList[i]) + ','

        resIpoDate = w.wss(allStockIdStr, 'ipo_date')

        stockIpoDict = {}
        for i in range(0, len(windStockIdList)) :
            aStockIdStr = str(windStockIdList[i])
            aDateStr = resIpoDate.Data[0][i].strftime('%Y-%m-%d')
            stockIpoDict[aDateStr] = aStockIdStr
        ipoDateDictList.append(stockIpoDict)
        collection_ipo_date.insert(ipoDateDictList)

        print('Finished writing base IPO date to database.')

        self.updateDataUpdateDate('ipo_date_update_date', ' ipo date ', currentDateStr = endDateStr)


    def updateIpoDateToMongoDB(self):


        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName

        updateDateStr = self.updateDateStr
        colName = self.colIpoDateName

        resStockID = w.wset('sectorconstituent', 'date=' + updateDateStr + ';sectorid=a001010100000000;field=wind_code')

        windStockIdList = resStockID.Data[0]

        allStockIdStr = ''

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]

        collection_ipo_date = db[colName]

        for i in range(0, len(windStockIdList)):
            allStockIdStr = allStockIdStr + str(windStockIdList[i]) + ','

        resIpoDate = w.wss(allStockIdStr, 'ipo_date')


        for i in range(0, len(windStockIdList)):
            aStockIdStr = str(windStockIdList[i])
            aDateStr = resIpoDate.Data[0][i].strftime('%Y%m%d')
            collection_ipo_date.update({}, {'$set': {aDateStr: aStockIdStr}}, upsert=True)


        print('IPO date for all stocks has been updated...')

        self.updateDataUpdateDate('ipo_date_update_date', ' IPO date for all stocks ')


    def writeBaseTradeStatusToMongoDB(self, endDateStr):
        startDateStr = self.baseStartDateStr
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colName = self.colTradeStatusName
        fillType = self.fillType

        stockExtractGrpNum = self.stockExtractGrpNum

        resStockID = w.wset('sectorconstituent', 'date=' + endDateStr + ';sectorid=a001010100000000;field=wind_code')

        stockIDdata = resStockID.Data[0]

        totalStockIdNum = len(stockIDdata)

        ###############################################################################################################
        #this shows stock number to retrieve in one group
        ###############################################################################################################

        splitStocksNum = math.floor(totalStockIdNum / stockExtractGrpNum)

        client = MongoClient(DBIP, DBPort)

        db = client[DBName]

        collection_trade_status = db[colName]

        for i in range(0, stockExtractGrpNum):
            stockIDOneLoopList = []
            for j in range(0, splitStocksNum):
                stockIDOneLoopList.append(i*splitStocksNum + j)
            oneStockSearchStr = ''
            for aStockID in stockIDOneLoopList:
                oneStockSearchStr = oneStockSearchStr + str(stockIDdata[aStockID]) + ', '
            resTradeStatus = w.wsd(oneStockSearchStr, 'trade_status', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)

            batchStockNum = len(resTradeStatus.Data)
            tradeStatusDataTime = resTradeStatus.Times

            for k in range(0, batchStockNum):
                aData = []
                aDateDict = {}
                aDateDict['StockID'] = stockIDdata[i * splitStocksNum + k]
                totalTransDaysNum = len(tradeStatusDataTime)
                for l in range(0, totalTransDaysNum):
                    aDateDict[tradeStatusDataTime[l].strftime('%Y%m%d')] = resTradeStatus.Data[k][l]
                aData.append(aDateDict)
                collection_trade_status.insert(aData)
            print('Finished writing trade status of stocks:' + oneStockSearchStr)

        ###############################################################################################################
        # After the looped retrieval, there could be still some stocks not retrieved, i.e., looped retrieved stocks may
        # not be the full set of available stocks, therefore we should locate the unretrieved stocks
        # by using mod function, we can check if remaining stocks is the full set of the available stock set
        # if not, we locate the remains
        ###############################################################################################################
        if (totalStockIdNum % stockExtractGrpNum) != 0:
            stockIDRemainsList = []
            for m in range((stockExtractGrpNum-1)*splitStocksNum + splitStocksNum, totalStockIdNum):
                stockIDRemainsList.append(m)
            oneStockSearchStr = ''

            for aStockID in stockIDRemainsList:
                oneStockSearchStr = oneStockSearchStr + str(stockIDdata[aStockID]) + ', '
            resTradeStatus = w.wsd(oneStockSearchStr, 'trade_status', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)
            batchStockNum = len(resTradeStatus.Data)
            tradeStatusDataTime = resTradeStatus.Times

            for k in range(0, batchStockNum):
                aData = []
                aDateDict = {}
                #aDateDict['StockID'] = stockIDRemainsList[k]
                aDateDict['StockID'] = stockIDdata[(stockExtractGrpNum-1)*splitStocksNum + splitStocksNum + k]
                for l in range(0, totalTransDaysNum):
                    aDateDict[tradeStatusDataTime[l].strftime('%Y%m%d')] = resTradeStatus.Data[k][l]
                aData.append(aDateDict)
                collection_trade_status.insert(aData)
            print('Finished writing remaining trade status of stocks:' + oneStockSearchStr)

        self.updateDataUpdateDate('trade_status_update_date', ' trade status ', currentDateStr = endDateStr)



    ###################################################################################################################
    ### Update Date Data as Current Date
    ###################################################################################################################
    def updateDataUpdateDate(self, strUpdatedKeyName, strUpdatedKeyAlias, currentDateStr = time.strftime('%Y-%m-%d')):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colSysData = self.colSysData

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_sys_data = db[colSysData]
        try:
            collection_sys_data.update({}, {'$set': {strUpdatedKeyName: currentDateStr}})
        except:
            print('Failed to update  date to MongoDB, please check.......')
            return()
        print('Successfully updated' + strUpdatedKeyAlias + 'into MongoDB.')


    def updateTradeStatusToMongoDB(self, startDateStr, endDateStr = time.strftime('%Y-%m-%d')):
        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colName = self.colTradeStatusName

        resStockID = w.wset('sectorconstituent', 'date=' + endDateStr + ';sectorid=a001010100000000;field=wind_code')

        stockIDdata = resStockID.Data[0]

        totalStockIdNum = len(stockIDdata)

        stockIdSearchStr = ''

        for n in range(0, totalStockIdNum):
            stockIdSearchStr = stockIdSearchStr + stockIDdata[n] + ', '
        resTradeStatus = w.wsd(stockIdSearchStr, 'trade_status', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)

        traeStatusTime = resTradeStatus.Times
        if (str(resTradeStatus.Data[0][0]).find('No data') != -1):
            print('Currently it is NOT in a transacation day and data is up to date, no more data to be updated by now...')
            return()

        totalTransDaysNum = len(traeStatusTime)

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_trade_status = db[colName]
        tradeStatusIdList = resTradeStatus.Data
        if (totalTransDaysNum >1 ):
            for n in range(0, totalStockIdNum):
                aStockIdStr = str(stockIDdata[n])
                for l in range(0, totalTransDaysNum):
                    aDateStr = traeStatusTime[l].strftime('%Y%m%d')
                    aTradeStatusStr = str(tradeStatusIdList[n][l])
                    collection_trade_status.update({'StockID': aStockIdStr}, {'$set': {aDateStr: aTradeStatusStr}}, upsert=True)
                    print('finished updating trade status of stock:' + str(stockIDdata[n]) + ', on the day of ' +  aDateStr)
        else:
            for n in range(0, totalStockIdNum):
                aStockIdStr = str(stockIDdata[n])
                for l in range(0, totalTransDaysNum):
                    aDateStr = traeStatusTime[l].strftime('%Y%m%d')
                    aTradeStatusStr = str(tradeStatusIdList[l][n])
                    collection_trade_status.update({'StockID': aStockIdStr}, {'$set': {aDateStr: aTradeStatusStr}}, upsert=True)
                    print('finished updating trade status of stock:' + str(stockIDdata[n]) + ', on the day of ' +  aDateStr)
        print('Finished updating trade status of all stocks of:' + stockIdSearchStr)

        self.updateDataUpdateDate('trade_status_update_date', ' trade status for all stocks ')



    ###################################################################################################################
    ### Obtain Data Update Date
    ###################################################################################################################
    def retrieveDataUpdateDate(self, strUpdateTimeKey):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colSysData = self.colSysData

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_sys_data = db[colSysData]


        sysDataRes = collection_sys_data.find({}, {strUpdateTimeKey: 1})

        dateListDict = sysDataRes[0]

        dataUpdateDate = dateListDict[strUpdateTimeKey]

        return (dataUpdateDate)

obj = StockGeneralData('2008-01-01')


baseEndDateStr = '2017-05-01'

#obj.writeBaseIpoDateToMongoDB(baseEndDateStr)
#obj.writeBaseTradeStatusToMongoDB(baseEndDateStr)


obj.updateIpoDateToMongoDB()
obj.updateTradeStatusToMongoDB(obj.retrieveDataUpdateDate('trade_status_update_date'))




